#include "rpc_base.hh"

#include "../../plugin/protocol/locator/rpc_path_req_generated.h"
#include "../../plugin/protocol/locator/rpc_path_ack_generated.h"
#include "../../plugin/protocol/locator/rpc_site_req_generated.h"
#include "../../plugin/protocol/locator/rpc_site_ack_generated.h"
#include "../../plugin/protocol/locator/rpc_entity_req_generated.h"
#include "../../plugin/protocol/locator/rpc_entity_ack_generated.h"

#include "../../corelib/network/thirdparty/knet/misc.h"
#include "../../corelib/util/os_util.hh"
#include "../../corelib/util/string_util.hh"
#include "../../plugin/base/flatbuffers_helper.hh"
#include "../../plugin/protocol/protocol_register.hh"
#include "../../plugin/base/entity.hh"

#ifdef WIN32
#include <WinSock2.h>
#pragma comment(lib, "wsock32.lib")
#else
#include <arpa/inet.h>
#endif

inline static uint64_t _htonll(uint64_t ui64) {
    return (((uint64_t)htonl((uint32_t)ui64)) << 32) + htonl(ui64 >> 32);
}

inline static uint64_t _ntohll(uint64_t ui64) {
    return (((uint64_t)ntohl((uint32_t)ui64)) << 32) + ntohl(ui64 >> 32);
}

void ProtocolRPCHeader::hton() {
    msgID = htonl(msgID);
    length = htonl(length);
    RPCHeaderLength = htonl(RPCHeaderLength);
}

void ProtocolRPCHeader::ntoh() {
    msgID = ntohl(msgID);
    length = ntohl(length);
    RPCHeaderLength = ntohl(RPCHeaderLength);
}

ProtocolRPCHeader::ProtocolRPCHeader() {
    length = 0;
    msgID = 0;
    RPCHeaderLength = 0;
}

ProtocolRPCHeader::ProtocolRPCHeader(const char * address) {
    length = *reinterpret_cast<const std::uint32_t*>(address);
    msgID = *reinterpret_cast<const std::uint32_t*>(address + sizeof(std::uint32_t));
    RPCHeaderLength = *reinterpret_cast<const std::uint32_t*>(address + sizeof(std::uint32_t) * 2);
    ntoh();
}

void ProtocolRPCHeader::toBuffer(char * address) {
    *reinterpret_cast<std::uint32_t*>(address) = htonl(length);
    *reinterpret_cast<std::uint32_t*>(address + sizeof(std::uint32_t)) = htonl(msgID);
    *reinterpret_cast<std::uint32_t*>(address + sizeof(std::uint32_t) * 2) = htonl(RPCHeaderLength);
}

RPC::RPC()
    : buffer_(new char[kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE]),
          skeleton_(nullptr), importer_(nullptr) {
}

RPC::~RPC() {
}

void RPC::registerSkeleton(RPCSkeletonBase * skeleton) {
    skeleton_ = skeleton;
}

bool RPC::buildMessage(std::uint32_t msgID, std::uint64_t RPCID, const google::protobuf::Message& msg,
    RPCCallBase* call, RPCTimeoutCB toCB, std::time_t timeout, std::size_t& length, std::uint64_t siteID, std::uint64_t& uuid) {
    static flatbuffers::FlatBufferBuilder builder(FLAT_BUFFER_SIZE);
    uuid = 0;
    if (!skeleton_ || !call) {
        return false;
    }
    builder.Clear();
    // Generates UUID of RPC
    auto UUID = ModuleRef.getUUID();
    if (pendingCallMap_.find(UUID) != pendingCallMap_.end()) {
        errorlog << "[RPC]Internal error:UUID explode" << endlog;
        return false;
    }
    if (!siteID) {
        auto RPCReq = CreateRPCPathReq(builder, UUID, RPCID);
        builder.Finish(RPCReq);
    } else {
        auto RPCReq = CreateRPCSiteReq(builder, siteID, UUID, RPCID);
        builder.Finish(RPCReq);
    }
    auto internalHeaderSize = builder.GetSize();
    auto result = msg.SerializeToArray(buffer_.get() + sizeof(ProtocolRPCHeader) + internalHeaderSize,
        kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE - internalHeaderSize - sizeof(ProtocolRPCHeader));
    if (!result) {
        errorlog << "[RPC]Internal error:msg.SerializeToArray, Message:" << msg.GetDescriptor()->name() << endlog;
        return false;
    }
    call->RPCID = RPCID;
    call->userTimeoutCB = toCB;
    call->timeout = static_cast<std::uint32_t>(timeout);
    call->timerCB = [&](TimerID, UserData UUID) {
        auto it = pendingCallMap_.find(UUID);
        if (it == pendingCallMap_.end()) {
            //errorlog << "[RPC]Internal error:Pending RPC not found, UUID:" << UUID << endlog;
            return;
        }
        auto call = it->second;
        if (call->userTimeoutCB) {
            call->userTimeoutCB();
        }
        skeleton_->dispose(call->RPCID, call);
        pendingCallMap_.erase(it);
    };
    // send ACK back to caller
    ProtocolRPCHeader header;
    header.msgID = msgID;
    header.length = sizeof(ProtocolRPCHeader) + internalHeaderSize + msg.ByteSize();
    header.RPCHeaderLength = internalHeaderSize;
    header.toBuffer(buffer_.get());
    memcpy(buffer_.get() + sizeof(ProtocolRPCHeader), builder.GetBufferPointer(), internalHeaderSize);
    length = header.length;
    call->timerID = ModuleRef.startTimerOnce(timeout, call->timerCB, UUID);
    if (call->timerID != 0) {
        pendingCallMap_[UUID] = call;
        uuid = UUID;
    } else {
        errorlog << "[RPC]Internal error:Start pending timer failed" << endlog;
        return false;
    }
    return true;
}

bool RPC::buildMessageEntity(std::uint32_t msgID, std::uint64_t RPCID, const google::protobuf::Message & msg,
    RPCCallBase * call, RPCTimeoutCB toCB, std::time_t timeout, std::size_t & length, std::uint64_t entityID,
    PathID& pathID, std::uint64_t& uuid) {
    static flatbuffers::FlatBufferBuilder builder(FLAT_BUFFER_SIZE);
    uuid = 0;
    if (!skeleton_ || !call) {
        return false;
    }
    builder.Clear();
    // Generates UUID of RPC
    auto UUID = ModuleRef.getUUID();
    if (pendingCallMap_.find(UUID) != pendingCallMap_.end()) {
        errorlog << "[RPC]Internal error:UUID explode" << endlog;
        return false;
    }
    pathID = ModuleRef.getEngine()->getLocator(entityID);
    if (pathID == INVALID_PATH_ID) {
        return false;
    }
    auto RPCReq = CreateRPCEntityReq(builder, entityID, pathID, UUID, RPCID);
    builder.Finish(RPCReq);
    auto internalHeaderSize = builder.GetSize();
    auto result = msg.SerializeToArray(buffer_.get() + sizeof(ProtocolRPCHeader) + internalHeaderSize,
        kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE - internalHeaderSize - sizeof(ProtocolRPCHeader));
    if (!result) {
        errorlog << "[RPC]Internal error:msg.SerializeToArray, Message:" << msg.GetDescriptor()->name() << endlog;
        return false;
    }
    call->RPCID = RPCID;
    call->userTimeoutCB = toCB;
    call->timeout = static_cast<std::uint32_t>(timeout);
    call->timerCB = [&](TimerID, UserData UUID) {
        auto it = pendingCallMap_.find(UUID);
        if (it == pendingCallMap_.end()) {
            //errorlog << "[RPC]Internal error:Pending RPC not found, UUID:" << UUID << endlog;
            return;
        }
        auto call = it->second;
        if (call->userTimeoutCB) {
            call->userTimeoutCB();
        }
        skeleton_->dispose(call->RPCID, call);
        pendingCallMap_.erase(it);
    };
    // send ACK back to caller
    ProtocolRPCHeader header;
    header.msgID = msgID;
    header.length = sizeof(ProtocolRPCHeader) + internalHeaderSize + msg.ByteSize();
    header.RPCHeaderLength = internalHeaderSize;
    header.toBuffer(buffer_.get());
    memcpy(buffer_.get() + sizeof(ProtocolRPCHeader), builder.GetBufferPointer(), internalHeaderSize);
    length = header.length;
    call->timerID = ModuleRef.startTimerOnce(timeout, call->timerCB, UUID);
    if (call->timerID != 0) {
        pendingCallMap_[UUID] = call;
        uuid = UUID;
    } else {
        errorlog << "[RPC]Internal error:Start pending timer failed" << endlog;
        return false;
    }
    return true;
}

bool RPC::RPCCallByPath(PathID pathID, std::uint64_t RPCID, const google::protobuf::Message & msg, RPCCallBase* call,
    RPCTimeoutCB toCB, std::time_t timeout) {
    std::size_t length = 0;
    std::uint64_t uuid = 0;
    if (buildMessage(static_cast<std::uint32_t>(LocatorMsgID::RPC_REQ_PATH), RPCID, msg, call, toCB, timeout, length, 0, uuid)) {
        // send to path peer
        if (!ModuleRef.getFramework()->talkPath(pathID, buffer_.get(), length)) {
            pendingCallMap_.erase(pendingCallMap_.find(uuid));
            return false;
        }
    } else {
        return false;
    }
    return true;
}

bool RPC::RPCCallBySite(const kratos::grid::SiteID & siteID, std::uint64_t RPCID, const google::protobuf::Message & msg,
    RPCCallBase * call, RPCTimeoutCB toCB, std::time_t timeout) {
    std::size_t length = 0;
    std::uint64_t uuid = 0;
    if (buildMessage(static_cast<std::uint32_t>(LocatorMsgID::RPC_REQ_SITE), RPCID, msg, call, toCB, timeout, length, siteID.u64(), uuid)) {
        // send to site peer
        if (!ModuleRef.getFramework()->talkSite(siteID, buffer_.get(), length)) {
            pendingCallMap_.erase(pendingCallMap_.find(uuid));
            return false;
        }
    } else {
        return false;
    }
    return true;
}

bool RPC::RPCCallByEntity(std::uint64_t entityID, std::uint64_t RPCID, const google::protobuf::Message & msg, RPCCallBase * call,
    RPCTimeoutCB toCB, std::time_t timeout) {
    std::size_t length = 0;
    PathID pathID = INVALID_PATH_ID;
    std::uint64_t uuid = 0;
    if (buildMessageEntity(static_cast<std::uint32_t>(LocatorMsgID::RPC_REQ_ENTITY), RPCID, msg, call, toCB, timeout, length, entityID,
        pathID, uuid)) {
        // send to path peer
        if (!ModuleRef.getFramework()->talkPath(pathID, buffer_.get(), length)) {
            pendingCallMap_.erase(pendingCallMap_.find(uuid));
            return false;
        }
    } else {
        return false;
    }
    return true;
}

bool RPC::loadPBProto(const std::string& path) {
    class MyMultiFileErrorCollector : public google::protobuf::compiler::MultiFileErrorCollector {
        void AddError(const string& fileName, int line, int column, const string& message) {
            errorlog << "Protobuf compiler error:(" << fileName << ":" << line << ":" << column << "), " << message << endlog;
        }
    };
    auto pbfilesPath = kratos::corelib::StringUtil::rtrim(path, "\\/");
    google::protobuf::compiler::DiskSourceTree sourceTree;
    sourceTree.MapPath("", path);
    importer_ = new google::protobuf::compiler::Importer(&sourceTree, NULL);
    std::vector<std::string> files;
    kratos::corelib::OS::getFiles(path, "proto", files);
    for (auto& file : files) {
        auto pbfilePath = kratos::corelib::StringUtil::getFileName(file) + ".proto";
        if (!importer_->Import(kratos::corelib::StringUtil::ltrim(pbfilePath, "\\/"))) {
            return false;
        }
    }
    return true;
}

google::protobuf::Message* RPC::newPBProtoMessage(const std::string& typeName) {
    auto it = messagePool_.find(typeName);
    if (it != messagePool_.end()) {
        if (!it->second.empty()) {
            auto msg = it->second.back();
            it->second.pop_back();
            msg->Clear();
            return msg;
        }
    }
    auto descriptor = importer_->pool()->FindMessageTypeByName(typeName);
    if (!descriptor) {
        return nullptr;
    }
    auto message = messageFactory_.GetPrototype(descriptor);
    if (!message) {
        return nullptr;
    }
    return message->New();
}

void RPC::deletePBProtoMessage(google::protobuf::Message* msg) {
    auto msgName = msg->GetDescriptor()->name();
    auto it = messagePool_.find(msgName);
    if (it == messagePool_.end()) {
        messagePool_[msgName].push_back(msg);
    } else {
        it->second.push_back(msg);
    }
}

void RPC::deletePBProtoMessage(google::protobuf::Message& msg) {
    auto msgName = msg.GetDescriptor()->name();
    auto it = messagePool_.find(msgName);
    if (it == messagePool_.end()) {
        messagePool_[msgName].push_back(&msg);
    } else {
        it->second.push_back(&msg);
    }
}

void RPC::onRPCPathMessage(PathID pathID, std::uint32_t msgID, const char * data, std::size_t size) {
    if (!skeleton_) {
        return;
    }
    ProtocolRPCHeader header(data);
    switch (static_cast<LocatorMsgID>(msgID)) {
        case LocatorMsgID::RPC_REQ_PATH:
            onPathReq(pathID, header, data, size);
            return;
        case LocatorMsgID::RPC_ACK_PATH:
            onPathAck(pathID, header, data, size);
            return;
    }
}

void RPC::onRPCSiteMessage(const kratos::grid::SiteID & siteID, std::uint32_t msgID, const char * data, std::size_t size) {
    if (!skeleton_) {
        return;
    }
    ProtocolRPCHeader header(data);
    switch (static_cast<LocatorMsgID>(msgID)) {
        case LocatorMsgID::RPC_REQ_SITE:
            onSiteReq(siteID, header, data, size);
            return;
        case LocatorMsgID::RPC_ACK_SITE:
            onSiteAck(siteID, header, data, size);
            return;
    }
}

void RPC::onRPCEntityMessage(PathID pathID, std::uint32_t msgID, const char * data, std::size_t size) {
    ProtocolRPCHeader header(data);
    switch (static_cast<LocatorMsgID>(msgID)) {
        case LocatorMsgID::RPC_REQ_ENTITY:
            onEntityReq(pathID, header, data, size);
            return;
        case LocatorMsgID::RPC_ACK_ENTITY:
            onEntityAck(pathID, header, data, size);
            return;
    }
}

void RPC::onPathReq(PathID pathID, const ProtocolRPCHeader& header, const char* data, std::size_t size) {
    if (!skeleton_) {
        return;
    }
    GET_FLATBUFFER_OBJ(RPCPathReq, req, data + sizeof(ProtocolRPCHeader));
    auto gap = sizeof(ProtocolRPCHeader) + header.RPCHeaderLength;
    try {
        skeleton_->onPathReq(pathID, req->RPCID(), req->UUID(), data + gap, size - gap);
    } catch (std::exception& e) {
        errorlog << "[RPC]RPC::onPathReq, Exception:" << e.what() << ", RPCID:" << req->RPCID() << endlog;
    }
}

void RPC::onPathAck(PathID, const ProtocolRPCHeader& header, const char* data, std::size_t size) {
    if (!skeleton_) {
        return;
    }
    GET_FLATBUFFER_OBJ(RPCPathAck, ack, data + sizeof(ProtocolRPCHeader));
    auto gap = sizeof(ProtocolRPCHeader) + header.RPCHeaderLength;
    auto it = pendingCallMap_.find(ack->UUID());
    if (it == pendingCallMap_.end()) {
        errorlog << "[RPC]RPC::onPathAck, UUID not found:" << ack->UUID() << ", RPCID:" << ack->RPCID() << endlog;
        return;
    }
    try {
        skeleton_->onPathAck(ack->RPCID(), it->second, data + gap, size - gap);
    } catch (std::exception& e) {
        errorlog << "[RPC]RPC::onPathAck, Exception:" << e.what() << ", RPCID:" << ack->RPCID() << endlog;
    }
    ModuleRef.cancelTimer(it->second->timerID);
    skeleton_->dispose(ack->RPCID(), it->second);
    pendingCallMap_.erase(it);
}

void RPC::onSiteReq(const kratos::grid::SiteID & siteID, const ProtocolRPCHeader & header,
    const char * data, std::size_t size) {
    if (!skeleton_) {
        return;
    }
    GET_FLATBUFFER_OBJ(RPCSiteReq, req, data + sizeof(ProtocolRPCHeader));
    auto gap = sizeof(ProtocolRPCHeader) + header.RPCHeaderLength;
    try {
        skeleton_->onSiteReq(siteID, req->RPCID(), req->UUID(), data + gap, size - gap);
    } catch (std::exception& e) {
        errorlog << "[RPC]RPC::onSiteReq, Exception:" << e.what() << ", RPCID:" << req->RPCID() << endlog;
    }
}

void RPC::onSiteAck(const kratos::grid::SiteID&, const ProtocolRPCHeader & header,
    const char * data, std::size_t size) {
    if (!skeleton_) {
        return;
    }
    GET_FLATBUFFER_OBJ(RPCSiteAck, ack, data + sizeof(ProtocolRPCHeader));
    auto gap = sizeof(ProtocolRPCHeader) + header.RPCHeaderLength;
    auto it = pendingCallMap_.find(ack->UUID());
    if (it == pendingCallMap_.end()) {
        errorlog << "[RPC]RPC::onSiteAck, UUID not found:" << ack->UUID() << ", RPCID:" << ack->RPCID() << endlog;
        return;
    }
    try {
        skeleton_->onSiteAck(ack->RPCID(), it->second, data + gap, size - gap);
    } catch (std::exception& e) {
        errorlog << "[RPC]RPC::onSiteAck, Exception:" << e.what() << ", RPCID:" << ack->RPCID() << endlog;
    }
    ModuleRef.cancelTimer(it->second->timerID);
    skeleton_->dispose(ack->RPCID(), it->second);
    pendingCallMap_.erase(it);
}

void RPC::onEntityReq(PathID pathID, const ProtocolRPCHeader & header, const char * data, std::size_t size) {
    GET_FLATBUFFER_OBJ(RPCEntityReq, req, data + sizeof(ProtocolRPCHeader));
    auto entity = EntityManagerRef.getEntity(req->entityID());
    if (entity) {
        if (!skeleton_) {
            return;
        }
        auto gap = sizeof(ProtocolRPCHeader) + header.RPCHeaderLength;
        try {
            skeleton_->onEntityReq(pathID, req->pathID(), req->entityID(), req->RPCID(), req->UUID(), data + gap, size - gap);
        } catch (std::exception& e) {
            errorlog << "[RPC]RPC::onEntityReq, Exception:" << e.what() << ", RPCID:" << req->RPCID() << endlog;
        }
    } else {
        auto targetPathID = ModuleRef.getEngine()->getEntityPathID(req->entityID());
        if (targetPathID != INVALID_PATH_ID) {
            ModuleRef.getEngine()->talkPath(targetPathID, data, size);
        }
    }
}

void RPC::onEntityAck(PathID, const ProtocolRPCHeader & header, const char * data, std::size_t size) {
    GET_FLATBUFFER_OBJ(RPCEntityAck, ack, data + sizeof(ProtocolRPCHeader));
    if (ModuleRef.getEngine()->isRouter()) {
        ModuleRef.getEngine()->talkPath(ack->pathID(), data, size);        
    } else {
        if (!skeleton_) {
            return;
        }
        auto gap = sizeof(ProtocolRPCHeader) + header.RPCHeaderLength;
        auto it = pendingCallMap_.find(ack->UUID());
        if (it == pendingCallMap_.end()) {
            errorlog << "[RPC]RPC::onEntityAck, UUID not found:" << ack->UUID() << ", RPCID:" << ack->RPCID() << endlog;
            return;
        }
        try {
            skeleton_->onEntityAck(ack->RPCID(), it->second, data + gap, size - gap);
        } catch (std::exception& e) {
            errorlog << "[RPC]RPC::onEntityAck, Exception:" << e.what() << ", RPCID:" << ack->RPCID() << endlog;
        }
        ModuleRef.cancelTimer(it->second->timerID);
        skeleton_->dispose(ack->RPCID(), it->second);
        pendingCallMap_.erase(it);
    }
}

void sendPath(PathID pathID, flatbuffers::FlatBufferBuilder& builder, ProtocolRPCHeader & header,
    const google::protobuf::Message & msg) {
    static char buffer[kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE] = { 0 };
    header.toBuffer(buffer);
    memcpy(buffer + sizeof(ProtocolRPCHeader), builder.GetBufferPointer(), builder.GetSize());
    msg.SerializeToArray(buffer + sizeof(ProtocolRPCHeader) + builder.GetSize(),
        sizeof(buffer) - builder.GetSize() - sizeof(ProtocolRPCHeader));
    // send to peer
    ModuleRef.getFramework()->talkPath(pathID, buffer, header.length);
}

void sendSite(const kratos::grid::SiteID & siteID, flatbuffers::FlatBufferBuilder& builder,
    ProtocolRPCHeader & header, const google::protobuf::Message & msg) {
    static char buffer[kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE] = { 0 };
    header.toBuffer(buffer);
    memcpy(buffer + sizeof(ProtocolRPCHeader), builder.GetBufferPointer(), builder.GetSize());
    msg.SerializeToArray(buffer + sizeof(ProtocolRPCHeader) + builder.GetSize(),
        sizeof(buffer) - builder.GetSize() - sizeof(ProtocolRPCHeader));
    // send to peer
    ModuleRef.getFramework()->talkSite(siteID, buffer, header.length);
}

void sendPathAck(PathID pathID, std::uint64_t UUID, std::uint64_t RPCID, const google::protobuf::Message& msg) {
    static flatbuffers::FlatBufferBuilder builder(kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE);
    builder.Clear();
    auto RPCAck = CreateRPCPathAck(builder, UUID, RPCID);
    builder.Finish(RPCAck);
    // send ACK back to caller
    ProtocolRPCHeader header;
    header.msgID = static_cast<std::uint32_t>(LocatorMsgID::RPC_ACK_PATH);
    header.length = sizeof(ProtocolRPCHeader) + builder.GetSize() + msg.ByteSize();
    header.RPCHeaderLength = builder.GetSize();
    sendPath(pathID, builder, header, msg);
}

void sendEntityAck(PathID pathID, PathID fromPathID, std::uint64_t UUID, std::uint64_t RPCID, const google::protobuf::Message & msg) {
    static flatbuffers::FlatBufferBuilder builder(kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE);
    builder.Clear();
    auto RPCAck = CreateRPCEntityAck(builder, fromPathID, UUID, RPCID);
    builder.Finish(RPCAck);
    // send ACK back to caller
    ProtocolRPCHeader header;
    header.msgID = static_cast<std::uint32_t>(LocatorMsgID::RPC_ACK_ENTITY);
    header.length = sizeof(ProtocolRPCHeader) + builder.GetSize() + msg.ByteSize();
    header.RPCHeaderLength = builder.GetSize();
    sendPath(pathID, builder, header, msg);
}

void sendSiteAck(const kratos::grid::SiteID & siteID, std::uint64_t UUID, std::uint64_t RPCID, const google::protobuf::Message& msg) {
    static flatbuffers::FlatBufferBuilder builder(kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE);
    builder.Clear();
    auto RPCAck = CreateRPCSiteAck(builder, siteID.u64(), UUID, RPCID);
    builder.Finish(RPCAck);
    // send ACK back to caller
    ProtocolRPCHeader header;
    header.msgID = static_cast<std::uint32_t>(LocatorMsgID::RPC_ACK_SITE);
    header.length = sizeof(ProtocolRPCHeader) + builder.GetSize() + msg.ByteSize();
    header.RPCHeaderLength = builder.GetSize();
    sendSite(siteID, builder, header, msg);
}
